"""
API速率限制中间件

功能：
- 基于IP地址的请求频率限制
- 支持不同端点的独立限流配置
- Redis后端存储，支持分布式部署
- 灵活的限流策略（固定窗口、滑动窗口）
"""

import time
from typing import Dict, Optional, Tuple

from fastapi import FastAPI, HTTPException, Request, status
from redis.asyncio import Redis
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response

from app.core import settings
from app.database.redis import get_redis


class RateLimitMiddleware(BaseHTTPMiddleware):
    """
    API速率限制中间件

    特性：
    - 基于Redis的分布式限流
    - 支持不同路径的不同限流策略
    - 友好的错误信息和头信息
    """

    def __init__(
        self,
        app,
        default_limit: int = 100,  # 默认每分钟100次
        default_window: int = 60,  # 默认时间窗口60秒
        path_specific_limits: Optional[Dict[str, Tuple[int, int]]] = None,
    ):
        super().__init__(app)
        self.default_limit = default_limit
        self.default_window = default_window
        self.path_specific_limits = path_specific_limits or {}

        # 重要的API端点限流配置
        self._important_endpoints = {
            "/api/v1/auth/login": (10, 60),  # 登录接口：每分钟10次
            "/api/v1/auth/register": (5, 60),  # 注册接口：每分钟5次
            "/api/v1/auth/captcha": (20, 60),  # 验证码接口：每分钟20次
        }
        # 初始化Redis客户端
        self._redis_client: Redis = get_redis()

    async def dispatch(self, request: Request, call_next) -> Response:
        """
        处理请求并检查速率限制

        Args:
            request: HTTP请求对象
            call_next: 调用链中的下一个处理函数

        Returns:
            Response: 处理后的HTTP响应

        Raises:
            HTTPException: 当超过速率限制时抛出429错误
        """
        # 跳过非API请求和OPTIONS请求
        if not request.url.path.startswith("/api/") or request.method == "OPTIONS":
            return await call_next(request)

        # 获取客户端标识（优先使用真实IP）
        client_id = self._get_client_identifier(request)

        # 获取当前端点的限流配置
        limit, window = self._get_rate_limit_config(request.url.path)

        # 检查速率限制
        if await self._is_rate_limited(client_id, request.url.path, limit, window):
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail=f"请求过于频繁，请{window}秒后再试。限制：{limit}次/{window}秒",
                headers={
                    "Retry-After": str(window),
                    "X-RateLimit-Limit": str(limit),
                    "X-RateLimit-Remaining": "0",
                    "X-RateLimit-Reset": str(int(time.time() + window)),
                },
            )

        # 处理请求并添加速率限制头信息
        response = await call_next(request)
        return self._add_rate_limit_headers(response, client_id, request.url.path, limit, window)

    def _get_client_identifier(self, request: Request) -> str:
        """
        获取客户端唯一标识

        Args:
            request: 请求对象

        Returns:
            str: 客户端标识（IP地址）
        """
        # 优先使用X-Forwarded-For头（代理后面）
        if "x-forwarded-for" in request.headers:
            ip = request.headers["x-forwarded-for"].split(",")[0].strip()
        else:
            ip = request.client.host if request.client else "unknown"

        return ip

    def _get_rate_limit_config(self, path: str) -> Tuple[int, int]:
        """
        获取指定路径的速率限制配置

        Args:
            path: 请求路径

        Returns:
            Tuple[int, int]: (限制次数, 时间窗口秒数)
        """
        # 检查重要端点配置
        for endpoint, config in self._important_endpoints.items():
            if path.startswith(endpoint):
                return config

        # 检查路径特定配置
        for pattern, config in self.path_specific_limits.items():
            if path.startswith(pattern):
                return config

        # 返回默认配置
        return self.default_limit, self.default_window

    async def _is_rate_limited(self, client_id: str, path: str, limit: int, window: int) -> bool:
        """
        检查客户端是否被限流

        Args:
            client_id: 客户端标识
            path: 请求路径
            limit: 限制次数
            window: 时间窗口

        Returns:
            bool: 是否被限流
        """
        if not self._redis_client:
            return False  # 如果没有Redis，不进行限流

        key = f"{settings.KEY_RATE_LIMITER}:{client_id}:{path}"

        try:
            # 使用Redis的INCR和EXPIRE实现固定窗口限流
            current = await self._redis_client.incr(key)
            if current == 1:
                await self._redis_client.expire(key, window)

            return current > limit
        except Exception:
            # Redis操作失败时，不进行限流
            return False

    async def _add_rate_limit_headers(
        self, response: Response, client_id: str, path: str, limit: int, window: int
    ) -> Response:
        """
        添加速率限制头信息到响应

        Args:
            response: 原始响应
            client_id: 客户端标识
            path: 请求路径
            limit: 限制次数
            window: 时间窗口

        Returns:
            Response: 添加了头信息的响应
        """
        if not self._redis_client:
            return response

        key = f"{settings.KEY_RATE_LIMITER}:{client_id}:{path}"

        try:
            current = await self._redis_client.get(key)
            remaining = max(0, limit - int(current or 0))
            reset_time = await self._redis_client.ttl(key)

            response.headers["X-RateLimit-Limit"] = str(limit)
            response.headers["X-RateLimit-Remaining"] = str(remaining)
            response.headers["X-RateLimit-Reset"] = str(reset_time)
        except Exception:
            pass  # 忽略Redis错误

        return response


def setup_rate_limit_middleware(app: FastAPI) -> None:
    """
    设置速率限制中间件

    Args:
        app: FastAPI应用实例
    """
    # 仅在非调试模式下启用速率限制
    if not settings.DEBUG:
        app.add_middleware(RateLimitMiddleware)
